package com.yanggu.flink.datastream_api.window

import cn.hutool.core.date.DateUtil
import com.yanggu.flink.datastream_api.pojo.{Event, UrlViewCount}
import com.yanggu.flink.datastream_api.source.ClickSource
import org.apache.flink.api.common.eventtime.{SerializableTimestampAssigner, WatermarkStrategy}
import org.apache.flink.api.common.state.ValueStateDescriptor
import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.api.scala.function.ProcessWindowFunction
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows
import org.apache.flink.streaming.api.windowing.time.Time
import org.apache.flink.streaming.api.windowing.triggers.{Trigger, TriggerResult}
import org.apache.flink.streaming.api.windowing.windows.TimeWindow
import org.apache.flink.util.Collector

import java.util.Date

object WindowTrigger {

  def main(args: Array[String]): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    env
      .addSource(new ClickSource())
      .assignTimestampsAndWatermarks(WatermarkStrategy
        .forMonotonousTimestamps()
        .withTimestampAssigner(new SerializableTimestampAssigner[Event] {
          override def extractTimestamp(element: Event, recordTimestamp: Long) = element.timestamp
        })
      )
      .keyBy(r => r.url)
      .window(TumblingEventTimeWindows.of(Time.seconds(10L)))
      .trigger(new MyTrigger)
      .process(new WindowTriggerResult)
      .print()

    env.execute()

  }

}


class MyTrigger extends Trigger[Event, TimeWindow] {

  private lazy val firstEventValueStateDescriptor = new ValueStateDescriptor[Boolean]("first-event", classOf[Boolean])

  override def onElement(event: Event,
                         time: Long,
                         timeWindow: TimeWindow,
                         triggerContext: Trigger.TriggerContext) = {
    val firstEvent = triggerContext.getPartitionedState(firstEventValueStateDescriptor)
    if (firstEvent.value()) {
      for (i <- timeWindow.getStart until timeWindow.getEnd by 1000) {
        triggerContext.registerEventTimeTimer(i)
      }
      firstEvent.update(true)
    }
    TriggerResult.CONTINUE
  }

  override def onProcessingTime(l: Long, w: TimeWindow, triggerContext: Trigger.TriggerContext) = TriggerResult.CONTINUE

  override def onEventTime(l: Long, w: TimeWindow, triggerContext: Trigger.TriggerContext) = TriggerResult.FIRE

  override def clear(w: TimeWindow, triggerContext: Trigger.TriggerContext) = {
    triggerContext.getPartitionedState(firstEventValueStateDescriptor).clear()
  }

}

class WindowTriggerResult extends ProcessWindowFunction[Event, UrlViewCount, String, TimeWindow] {

  override def process(key: String, context: Context, elements: Iterable[Event], out: Collector[UrlViewCount]): Unit = {
    out.collect(UrlViewCount(key, elements.size, DateUtil.formatDateTime(new Date(context.window.getStart)), DateUtil.formatDateTime(new Date(context.window.getEnd))))
  }

}
